#!/usr/bin/env python3
"""
背景移除命令行工具 - 基于 rembg 的批量图片处理工具
"""

import os
import sys
import argparse
from pathlib import Path
import logging
from tqdm import tqdm

# 默认模型路径
DEFAULT_MODEL_DIR = Path(os.getcwd()) / "models"

try:
    from rembg import remove, new_session
except ImportError:
    print("错误：请先安装rembg: pip install rembg")
    sys.exit(1)

# 支持的输入图片格式
SUPPORTED_INPUT_FORMATS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp'}
# 输出格式（固定为PNG）
OUTPUT_FORMAT = '.png'

def setup_logging(verbose=False):
    """设置日志级别"""
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format='%(asctime)s - %(levelname)s - %(message)s',
        datefmt='%H:%M:%S'
    )

def setup_model_dir(model_dir):
    """设置模型目录并确保其存在"""
    model_dir_path = Path(model_dir)
    model_dir_path.mkdir(parents=True, exist_ok=True)
    
    # 设置环境变量，让 rembg 使用指定的模型目录
    os.environ['U2NET_HOME'] = str(model_dir_path)
    
    logging.info(f"模型目录设置为: {model_dir_path}")
    return model_dir_path

def get_session(model_name):
    """创建模型会话"""
    try:
        return new_session(model_name)
    except Exception as e:
        logging.error(f"创建模型会话失败: {e}")
        # 检查是否是模型文件缺失
        if "not found" in str(e).lower() or "download" in str(e).lower():
            logging.info(f"请确保模型 {model_name} 已下载到模型目录")
            logging.info(f"模型目录: {os.environ.get('U2NET_HOME', '未设置')}")
            logging.info("首次使用需要下载模型，这可能需要一些时间...")
        raise

def process_single_image(input_path, output_path, session, args):
    """处理单张图片"""
    try:
        with open(input_path, 'rb') as i:
            input_data = i.read()
            
            # 准备remove函数的参数
            remove_kwargs = {
                'session': session,
                'post_process_mask': args.post_process,
                'only_mask': args.only_mask,
            }
            
            # 如果启用alpha matting，添加相关参数
            if args.alpha_matting:
                remove_kwargs.update({
                    'alpha_matting': True,
                    'alpha_matting_foreground_threshold': args.foreground_threshold,
                    'alpha_matting_background_threshold': args.background_threshold,
                    'alpha_matting_erode_size': args.erode_size,
                })
            
            # 如果指定了背景颜色
            if args.bgcolor:
                remove_kwargs['bgcolor'] = args.bgcolor
            
            # 处理图片
            output_data = remove(input_data, **remove_kwargs)
            
            # 确保输出目录存在
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            with open(output_path, 'wb') as o:
                o.write(output_data)
                
        return True
        
    except Exception as e:
        logging.error(f"处理图片 {input_path} 失败: {e}")
        return False

def get_image_files(input_path):
    """获取所有图片文件（不递归子目录）"""
    input_path = Path(input_path)
    image_files = []
    
    if input_path.is_file():
        # 单个文件
        if input_path.suffix.lower() in SUPPORTED_INPUT_FORMATS:
            return [input_path]
        else:
            return []
    
    # 目录 - 仅查找当前目录下的图片文件，不递归子目录
    for file_path in input_path.iterdir():
        if file_path.is_file() and file_path.suffix.lower() in SUPPORTED_INPUT_FORMATS:
            image_files.append(file_path)
    
    return image_files

def get_output_path(input_path, output_path, is_batch=False):
    """生成输出路径，确保使用PNG格式"""
    input_path = Path(input_path)
    
    if output_path:
        output_path_obj = Path(output_path)
        
        # 判断输出路径是文件还是目录
        if output_path_obj.suffix.lower() in SUPPORTED_INPUT_FORMATS and not is_batch:
            # 单个文件输出，强制使用PNG格式
            return output_path_obj.parent / f"{output_path_obj.stem}{OUTPUT_FORMAT}"
        else:
            # 目录输出
            if is_batch:
                # 批量处理：在输出目录中保持原文件名但改为PNG格式
                return output_path_obj / f"{input_path.stem}{OUTPUT_FORMAT}"
            else:
                # 单个文件处理到目录
                return output_path_obj / f"{input_path.stem}{OUTPUT_FORMAT}"
    else:
        # 默认输出到原目录，添加_nobg后缀并使用PNG格式
        return input_path.parent / f"{input_path.stem}_nobg{OUTPUT_FORMAT}"

def process_batch(input_path, output_path, session, args):
    """批量处理图片"""
    input_path = Path(input_path)
    
    # 收集所有图片文件
    image_files = get_image_files(input_path)
    
    if not image_files:
        logging.error("未找到支持的图片文件")
        return False
    
    logging.info(f"找到 {len(image_files)} 个图片文件")
    logging.info(f"输出格式: {OUTPUT_FORMAT}")
    
    # 如果指定了输出路径，提前创建输出目录
    if output_path:
        output_path_obj = Path(output_path)
        # 判断输出路径是文件还是目录
        if output_path_obj.suffix.lower() in SUPPORTED_INPUT_FORMATS and len(image_files) == 1:
            # 单个文件输出，创建父目录
            output_path_obj.parent.mkdir(parents=True, exist_ok=True)
        else:
            # 目录输出，创建目录
            output_path_obj.mkdir(parents=True, exist_ok=True)
    
    # 处理进度条
    success_count = 0
    with tqdm(total=len(image_files), desc="处理进度") as pbar:
        for img_path in image_files:
            # 计算输出路径（强制使用PNG格式）
            is_batch = len(image_files) > 1
            out_file = get_output_path(img_path, output_path, is_batch)
            
            # 确保输出目录存在
            out_file.parent.mkdir(parents=True, exist_ok=True)
            
            # 处理图片
            if process_single_image(img_path, out_file, session, args):
                success_count += 1
                pbar.set_postfix({
                    "状态": "成功", 
                    "文件": img_path.name[:20] + "..." if len(img_path.name) > 20 else img_path.name,
                    "输出": out_file.name
                })
            else:
                pbar.set_postfix({
                    "状态": "失败", 
                    "文件": img_path.name[:20] + "..." if len(img_path.name) > 20 else img_path.name
                })
            
            pbar.update(1)
    
    logging.info(f"处理完成: {success_count}/{len(image_files)} 个文件成功")
    return success_count > 0

def parse_bgcolor(color_str):
    """解析背景颜色参数"""
    if not color_str:
        return None
    
    try:
        # 支持格式: "255,255,255" 或 "255,255,255,128"
        colors = [int(x.strip()) for x in color_str.split(',')]
        
        if len(colors) == 3:
            return (colors[0], colors[1], colors[2], 255)  # 默认不透明
        elif len(colors) == 4:
            return tuple(colors)
        else:
            raise ValueError("颜色参数格式错误")
    except Exception as e:
        logging.error(f"背景颜色解析错误: {e}")
        return None

def main():
    parser = argparse.ArgumentParser(
        description="背景移除工具 - 批量处理图片去除背景（默认输出PNG格式）",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 处理单张图片（默认输出PNG格式）
  python bg_remover.py input.jpg -o output.png
  
  # 批量处理目录（所有输出均为PNG格式）
  python bg_remover.py ./input_images -o ./output_images
  
  # 使用人像分割模型
  python bg_remover.py input.jpg -m u2net_human_seg --post-process
  
  # 启用alpha matting处理发丝细节
  python bg_remover.py input.jpg --alpha-matting --erode-size 15
  
  # 设置白色背景
  python bg_remover.py input.jpg --bgcolor "255,255,255"
  
  # 只生成掩码
  python bg_remover.py input.jpg --only-mask
  
  # 指定模型目录
  python bg_remover.py input.jpg --model-dir /path/to/models

注意:
  所有输出图片均为PNG格式，以保持透明背景质量

模型说明:
  u2net: 通用模型，平衡效果和速度
  u2net_human_seg: 人像分割专用
  isnet-general-use: 最高质量，速度较慢
  birefnet-general: 通用模型，模型较大
        """
    )
    
    # 输入输出参数
    parser.add_argument('input', help='输入文件或目录路径')
    parser.add_argument('-o', '--output', help='输出文件或目录路径（默认输出PNG格式）')
    
    # 模型参数
    parser.add_argument('-m', '--model', 
                       choices=["u2net",
                                "u2netp",
                                "u2net_human_seg",
                                "u2net_cloth_seg",
                                "silueta",
                                "isnet-general-use",
                                "isnet-anime",
                                "sam",
                                "birefnet-general", 
                                "birefnet-general-lite",
                                "birefnet-portrait",
                                "birefnet-dis",
                                "birefnet-hrsod",
                                "birefnet-cod",
                                "birefnet-massive"],
                       default='u2net',
                       help='选择模型 (默认: u2net)')
    
    # 新增：模型目录参数
    parser.add_argument('--model-dir', 
                       default=None,
                       help=f'模型存储目录 (默认: {DEFAULT_MODEL_DIR})')
    
    # 处理参数
    parser.add_argument('--post-process', action='store_true',
                       help='对掩码进行后处理（平滑和边缘优化）')
    parser.add_argument('--only-mask', action='store_true',
                       help='只返回黑白掩码')
    
    # Alpha matting 参数
    parser.add_argument('--alpha-matting', action='store_true',
                       help='启用alpha matting获得更精细边缘')
    parser.add_argument('--foreground-threshold', type=int, default=240,
                       help='前景阈值 (默认: 240)')
    parser.add_argument('--background-threshold', type=int, default=10,
                       help='背景阈值 (默认: 10)')
    parser.add_argument('--erode-size', type=int, default=10,
                       help='腐蚀结构大小 (默认: 10)')
    
    # 背景颜色
    parser.add_argument('--bgcolor', 
                       help='背景颜色，格式: "R,G,B" 或 "R,G,B,A" (例如: "255,255,255" 或 "0,0,255,128")')
    
    # 其他参数
    parser.add_argument('-v', '--verbose', action='store_true',
                       help='显示详细日志')
    
    args = parser.parse_args()
    
    # 设置日志
    setup_logging(args.verbose)
    
    # 设置模型目录
    setup_model_dir(args.model_dir or DEFAULT_MODEL_DIR)
    
    # 解析背景颜色
    if args.bgcolor:
        args.bgcolor = parse_bgcolor(args.bgcolor)
        if args.bgcolor is None:
            return
    
    # 验证输入路径
    input_path = Path(args.input)
    if not input_path.exists():
        logging.error(f"输入路径不存在: {args.input}")
        return
    
    # 处理输出路径
    output_path = None
    if args.output:
        output_path = Path(args.output)
    
    # 检查单文件输出情况
    input_image_files = get_image_files(input_path)
    if len(input_image_files) > 1 and output_path and output_path.exists() and output_path.is_file():
        logging.error("多个输入文件但输出是已存在的文件，请指定输出目录")
        return
    
    try:
        # 创建模型会话
        logging.info(f"加载模型: {args.model}")
        session = get_session(args.model)
        
        # 显示处理参数
        logging.info("处理参数:")
        logging.info(f"  输出格式: {OUTPUT_FORMAT}")
        logging.info(f"  后处理: {args.post_process}")
        logging.info(f"  只生成掩码: {args.only_mask}")
        logging.info(f"  Alpha Matting: {args.alpha_matting}")
        if args.alpha_matting:
            logging.info(f"    前景阈值: {args.foreground_threshold}")
            logging.info(f"    背景阈值: {args.background_threshold}")
            logging.info(f"    腐蚀大小: {args.erode_size}")
        if args.bgcolor:
            logging.info(f"  背景颜色: {args.bgcolor}")
        
        # 开始处理
        success = process_batch(args.input, output_path, session, args)
        
        if success:
            logging.info("批量处理完成！")
        else:
            logging.error("处理过程中出现错误")
            sys.exit(1)
            
    except KeyboardInterrupt:
        logging.info("用户中断处理")
    except Exception as e:
        logging.error(f"处理失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()